import os
import subprocess

from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

this_dir = os.path.dirname(os.path.abspath(__file__))

def write_version_to_file(version, target_file):
    with open(target_file, 'w') as f:
        print('__version__ = "%s"' % version, file=f)

if __name__ == '__main__':
    version = '0.1.0'
    write_version_to_file(version, './version.py')
    os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;8.0"
    setup(
        name='grasp_nms',
        version=version,
        description='grasp distance metric by grasp_center and approach_vector and closing_vector and grasp_width',
        install_requires=[
            'numpy',
            'torch>=1.1',
        ],
        author='Wei wei',
        author_email='wei.wei2018@ia.ac.cn',
        license='Apache License 2.0',
        packages=find_packages(),
        cmdclass={'build_ext': BuildExtension},
        ext_modules=[
            CUDAExtension(
                name='grasp_cuda',
                #module='pcdet.ops.iou3d_nms',
                sources=[
                    'src/grasp_api.cpp',
                    'src/grasp.cpp',
                    'src/roipoint_pool3d.cpp',
                    'src/grasp_nms_kernel.cu',
                    'src/roipoint_pool3d_kernel.cu'
                ],
            ),
        ],
    )
